#ifndef THREAD_POOL_H
#define THREAD_POOL_H

#include <vector>
#include <queue>
#include <memory>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <future>
#include <functional>
#include <stdexcept>
#include <atomic>

namespace common {
namespace thread {

class ThreadPool {
public:
    // 创建线程池，指定线程数量
    ThreadPool(size_t threads);
    
    // 提交任务到线程池，返回future
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args) 
        -> std::future<typename std::result_of<F(Args...)>::type>;
    
    // 获取工作线程数量
    size_t getThreadCount() const { return workers_.size(); }
    
    // 获取当前等待队列大小
    size_t getQueueSize() const;
    
    // 获取已完成任务数
    uint64_t getCompletedTaskCount() const { return completedTasks_; }
    
    // 析构函数
    ~ThreadPool();
    
private:
    // 工作线程向量
    std::vector<std::thread> workers_;
    // 任务队列
    std::queue<std::function<void()>> tasks_;
    
    // 同步
    mutable std::mutex queue_mutex_;
    std::condition_variable condition_;
    
    // 线程池控制标志
    std::atomic<bool> stop_;
    
    // 已完成任务数
    std::atomic<uint64_t> completedTasks_;
};

// 线程池实现
inline ThreadPool::ThreadPool(size_t threads)
    : stop_(false), completedTasks_(0) {
    for(size_t i = 0; i < threads; ++i)
        workers_.emplace_back(
            [this] {
                while(true) {
                    std::function<void()> task;
                    
                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex_);
                        this->condition_.wait(lock,
                            [this] { return this->stop_ || !this->tasks_.empty(); });
                        
                        if(this->stop_ && this->tasks_.empty())
                            return;
                        
                        task = std::move(this->tasks_.front());
                        this->tasks_.pop();
                    }
                    
                    task();
                    ++completedTasks_;
                }
            }
        );
}

// 添加新任务到线程池
template<class F, class... Args>
auto ThreadPool::enqueue(F&& f, Args&&... args) 
    -> std::future<typename std::result_of<F(Args...)>::type> {
    using return_type = typename std::result_of<F(Args...)>::type;

    auto task = std::make_shared<std::packaged_task<return_type()>>(
        std::bind(std::forward<F>(f), std::forward<Args>(args)...)
    );
        
    std::future<return_type> res = task->get_future();
    {
        std::unique_lock<std::mutex> lock(queue_mutex_);

        // 不接受新任务如果线程池已停止
        if(stop_)
            throw std::runtime_error("enqueue on stopped ThreadPool");

        tasks_.emplace([task](){ (*task)(); });
    }
    condition_.notify_one();
    return res;
}

// 获取当前等待队列大小
inline size_t ThreadPool::getQueueSize() const {
    std::unique_lock<std::mutex> lock(queue_mutex_);
    return tasks_.size();
}

// 停止线程池
inline ThreadPool::~ThreadPool() {
    {
        std::unique_lock<std::mutex> lock(queue_mutex_);
        stop_ = true;
    }
    condition_.notify_all();
    
    for(std::thread &worker: workers_)
        worker.join();
}

} // namespace thread
} // namespace common

#endif // THREAD_POOL_H 